package org.example.algorithm.dp;

public class CountSquaresSolution {

    public static void main(String[] args) {
        CountSquaresSolution solution = new CountSquaresSolution();
        int[][] matrix = {
                {0, 0, 0},
                {0, 1, 0},
                {0, 1, 0},
                {1, 1, 1},
                {1, 1, 0}
        };
        int res = solution.countSquares3(matrix);
        System.out.println(res);
    }

    //dp[i+1][j+1]代表以位置(i,j)为右下角矩形的最大边长
    //dp[i+1][j+1] = min(dp[i][j+1], dp[i+1][j], dp[i][j]) + 1
    public int countSquares(int[][] matrix) {
        int row = matrix.length;
        int col = matrix[0].length;
        int[][] dp = new int[row+1][col+1];
        int res = 0;
        for (int i=0;i<row;i++) {
            for (int j=0;j<col;j++) {
                if (matrix[i][j] == 1) {
                    dp[i+1][j+1] = Math.min(Math.min(dp[i][j+1], dp[i+1][j]), dp[i][j]) + 1;
                    res += dp[i+1][j+1];
                }
            }
        }
        return res;
    }

    //暴力破解
    public int countSquares2(int[][] matrix) {
        int row = matrix.length;
        int col = matrix[0].length;
        int res = 0;
        for (int i=0;i<row;i++) {
            for (int j=0; j<col; j++) {
                if (matrix[i][j] == 1) {
                    res++;
                    int maxSide = Math.min(row-i, col-j);
                    boolean flag = true;
                    for (int k=1;k<maxSide;k++) {
                        for (int m=0;m<=k;m++) {
                            if (matrix[i+k][j+m] == 0 || matrix[i+m][j+k] == 0) {
                                flag = false;
                                break;
                            }
                        }
                        if (!flag) {
                            break;
                        }
                        res++;
                    }
                }
            }
        }
        return res;
    }

    //前缀和+二分查找
    public int countSquares3(int[][] matrix) {
        int row = matrix.length;
        int col = matrix[0].length;
        int[][] sum = new int[row+1][col+1];
        for (int i=0;i<row;i++) {
            for (int j=0;j<col;j++) {
                sum[i+1][j+1] = sum[i+1][j] + sum[i][j+1] - sum[i][j] + matrix[i][j];
            }
        }
        int res = 0;
        for (int i=0;i<row;i++) {
            for (int j=0;j<col;j++) {
                int l = 1;
                int r = Math.min(row-i, col-j);
                int maxSide = 0;
                while (l <= r) {
                    int mid = (l + r) / 2;
                    if (checkSquareSum(sum, i, j, mid)) {
                        maxSide = mid;
                        l = mid + 1;
                    } else {
                        r = mid - 1;
                    }
                }
                res += maxSide;
            }
        }
        return res;
    }

    private boolean checkSquareSum(int[][] sum, int i, int j, int mid) {
        //左上角(i,j)，右下角(i+mid-1, j+mid-1)
        return sum[i + mid][j + mid] - sum[i + mid][j] - sum[i][j + mid] + sum[i][j] == mid * mid;
    }
}
